# encoding=utf-8
import shutil
import os
from skimage import io
import SimpleITK as sitk
import numpy as np
import traceback
import nibabel as nib


class NnunetPreProcess:
    # 封装nnunet的数据处理工具函数

    @staticmethod
    def convert2png(npy_path: str, axis: str = 'z') -> np.array:
        npy_data = np.asanyarray(nib.load(npy_path).dataobj)
        npy_data = np.rot90(np.fliplr(npy_data))
        if axis == 'x':
            return npy_data[0, :, :]
        elif axis == 'y':
            return npy_data[:, 0, :]
        else:
            return npy_data[:, :, 0]

    @staticmethod
    def convert2nifti(mat_path: str, save_path: str, spacing=(999, 1, 1)):
        """
            框架file_conversions的convert_2d_image_to_nifti
            简化执行思路的二次封装版本

            mat_path是图像的绝对路径,
            save_path是存放图片的路径
        """
        try:
            file_name = os.path.basename(mat_path)
            stem, suffix = os.path.splitext(file_name)
            save_name_list = []
            # abs_save_path = os.path.join(save_path, stem + ".nii.gz")

            img = io.imread(mat_path)
            if len(img.shape) == 2:  # 2d image with no color channels
                img = img[None, None]  # add dimensions
            else:
                assert len(img.shape) == 3, "image should be 3d with color channel last but has shape %s" % str(
                    img.shape)
                img = img.transpose((2, 0, 1))
                img = img[:, None]

            # 这里拆分三个通道,分别进行保存
            for idx, e_img in enumerate(img):
                e_img = e_img.astype(np.uint32)

                itk_img = sitk.GetImageFromArray(e_img)
                itk_img.SetSpacing(list(spacing)[::-1])

                nii_name = "{}_{}.nii.gz".format(stem, idx)
                save_name_list.append(nii_name)

                sitk.WriteImage(itk_img, os.path.join(save_path, nii_name))

            return save_name_list
        except Exception as e:
            print(traceback.format_exc())
            return None

    @staticmethod
    def try_rm_dir(path: str):
        try:
            shutil.rmtree(path)
            return True
        except Exception as e:
            print(traceback.format_exc())
            return False
